function [D,invtausq_d,P,llik_trace] = update_covariate_model(B,C,H,BURNIN,NEPOCH,D,invtausq_d)
%
% [D,invTauSq_d,P,llik_trace] = update_covariate_model(B,C,H,BURNIN,NEPOCH, initial_D, initial_invtausq_d)
%
% Inference of covariate model
%
% B|H,C,D ~ logit(H+C*D') and D ~ group lasso
%
% B          = ind x probe observed methylation beta matrix [0, 1]
% C          = ind x covariate matrix
% H          = ind x probe hidden activity
% BURNIN     = burn-in optimization time without TauSq
% NEPOCH     = optimization time with TauSq
% D          = probe x (covariate+1) where constant column appended in the end
% invTauSq_d = shrinkage of each covariate (including intercept)
% P          = predicted activity according to the logit model
%
% code: Yongjin Park, ypp@csail.mit.edu
%

    TOL = 1e-3;

    % ================================================================
    [Nind, Nprobe] = size(B);
    [nind, Ncov] = size(C);
    assert(nind == Nind);

    C = [C,ones(Nind,1,'single')];
    Ncov = Ncov + 1;
    [nind, nprobe] = size(H);
    assert(Nind == nind && nprobe == Nprobe);

    % ================================================================
    if nargin > 5,
        assert(size(D,1) == Nprobe && size(D,2) == Ncov);
        assert(size(invtausq_d,1) == Ncov && size(invtausq_d,2) == 1);
    else
        D = zeros(Nprobe,Ncov,'single');
        invtausq_d = 0.01*Nind*ones(Ncov,1,'single');
    end

    % ================================================================
    llik_trace = NaN(BURNIN+NEPOCH,1,'single');

    % ================================================================
    % take care of missing values
    % (1) W(i,j) = 0 if B(i,j) missing
    % (2) when optimizing D(:,k), W(i,:) = 0 if C(i,k) missing
    missing_probe = isnan(B);
    missing_covar = isnan(C);

    % simply pad zeros for missing spots
    B(missing_probe) = 0;
    C(missing_covar) = 0;

    for iter = 1:(BURNIN + NEPOCH),

        % 1. construct quadratic approximation
        F = C*D';
        P = 1./(1+exp(-F-H));

        P(P < 1e-4) = 1e-4;
        P(P > 1 - 1e-4) = 1 - 1e-4;

        W = P.*(1-P);
        X = F + (B - P) ./ W;
        
        W(missing_probe) = 0;

        % 2. update each coordinate
        R = X - F;
        WCsq = W'*(C.^2);
        
        for k = 1:Ncov,
            R_exc = R + bsxfun(@times, C(:,k), D(:,k).');
            denom = invtausq_d(k) + WCsq(:,k);
            dd = (W.*R_exc)'*C(:,k) ./ denom;

            D(:,k) = dd;
            R = R_exc - bsxfun(@times, C(:,k), D(:,k).');
        end

        F = C*D';

        llik = sum(sum(B.*F - log(1+exp(F))));
        llik_trace(iter) = llik;

        if iter > BURNIN,
            invtausq_d = arrayfun(@(x) max(0.1/Nprobe,min(10*Nprobe,x)), 1./mean(D.^2)');
        end

        if iter > BURNIN + 10,
            llik_prev = mean(llik_trace((iter-10):(iter-6)));
            llik_curr = mean(llik_trace((iter-5):(iter-1)));

            if abs(llik_prev - llik_curr)/abs(TOL + llik_curr) < TOL,
                llik_trace = llik_trace(1:iter);
                fprintf(2,'Converged\r');
                break;
            end
        end
        
        fprintf(2,'Iter = %03d, LLIK = %.4e\r',iter,llik);

    end

end
